# -*- coding: utf-8 -*-
"""
Created on Wed Aug 18 09:39:30 2021

@author: MLoong
"""

import torch.nn as nn
import torchvision.models as models
import torch

class Model(nn.Module):
    """
    simply create a 2-branch network, and concat global pooled feature vector.
    each branch = single resnet34
    """
    def __init__(self,classes=1,pretrained=True):
        super(Model, self).__init__()
        self.fundus_branch =  nn.Sequential(*list(models.resnet50(pretrained=pretrained).children())[:-1])
        self.oct_branch = nn.Sequential(*list(models.resnet50(pretrained=pretrained).children())[:-1])
        self.oct_branch[0] = nn.Conv2d(256, 64,kernel_size=7,stride=2,padding=3)
        self.fc = classifier_layer(classes=3)
        

    def forward(self, fundus_img, oct_img):
        b1 = self.fundus_branch(fundus_img)
        b2 = self.oct_branch(oct_img)
        cobine_feture=torch.cat([b1, b2], 1)
        logit=self.fc(cobine_feture)
        # print(torch.add(b1,0.7,0))
        # print(torch.add(b2,0.3,0))
        # logit=torch.add(torch.add(b1,0.5,0),torch.add(b2,0.5,0))
        # print(logit)
        return logit
    
def classifier_layer(classes=3):
    fc= nn.Sequential(nn.Flatten(),
                      # nn.Linear(2048, 64), 
                      # # nn.ReLU(),  
                      # nn.Dropout(0.50), 
                      nn.Linear(4096,classes),
                      nn.Softmax(dim=1))
    return fc
if __name__=="__main__":
    model=Model()
    print(model)
#    print(*list(model.children()))